/////////////////////////////////////////////////////////////////
// SparseMatrix.h
//
// Sparse matrix computations
/////////////////////////////////////////////////////////////////

#ifndef SPARSEMATRIX_H
#define SPARSEMATRIX_H

#include <iostream>

using namespace std;

const float POSTERIOR_CUTOFF = 0.01;         // minimum posterior probability
// value that is maintained in the
// sparse matrix representation

typedef pair<int, float> PIF;                 // Sparse matrix entry type
//   first --> column
//   second --> value

/////////////////////////////////////////////////////////////////
// SparseMatrix
//
// Class for sparse matrix computations
/////////////////////////////////////////////////////////////////

class SparseMatrix {

	int seq1Length, seq2Length;                     // dimensions of matrix
	VI rowSize;                              // rowSize[i] = # of cells in row i
	SafeVector<PIF> data;                           // data values
	SafeVector<SafeVector<PIF>::iterator> rowPtrs; // pointers to the beginning of each row

	/////////////////////////////////////////////////////////////////
	// SparseMatrix::SparseMatrix()
	//
	// Private constructor.
	/////////////////////////////////////////////////////////////////

	SparseMatrix() {
	}

public:

	/////////////////////////////////////////////////////////////////
	// SparseMatrix::SparseMatrix()
	//
	// Constructor.  Builds a sparse matrix from a posterior matrix.
	// Note that the expected format for the posterior matrix is as
	// a (seq1Length+1) x (seq2Length+1) matrix where the 0th row
	// and 0th column are ignored (they should contain all zeroes).
	/////////////////////////////////////////////////////////////////

	SparseMatrix(int seq1Length, int seq2Length, const VF &posterior) :
			seq1Length(seq1Length), seq2Length(seq2Length) {

		int numCells = 0;

		assert(seq1Length > 0);
		assert(seq2Length > 0);

		// calculate memory required; count the number of cells in the
		// posterior matrix above the threshold
		VF::const_iterator postPtr = posterior.begin();
		for (int i = 0; i <= seq1Length; i++) {
			for (int j = 0; j <= seq2Length; j++) {
				if (*(postPtr++) >= POSTERIOR_CUTOFF) {
					assert(i != 0 && j != 0);
					numCells++;
				}
			}
		}

		// allocate memory
		data.resize(numCells);
		rowSize.resize(seq1Length + 1);
		rowSize[0] = -1;
		rowPtrs.resize(seq1Length + 1);
		rowPtrs[0] = data.end();

		// build sparse matrix
		postPtr = posterior.begin() + seq2Length + 1; // note that we're skipping the first row here
		SafeVector<PIF>::iterator dataPtr = data.begin();
		for (int i = 1; i <= seq1Length; i++) {
			postPtr++;              // and skipping the first column of each row
			rowPtrs[i] = dataPtr;
			for (int j = 1; j <= seq2Length; j++) {
				if (*postPtr >= POSTERIOR_CUTOFF) {
					dataPtr->first = j;
					dataPtr->second = *postPtr;
					dataPtr++;
				}
				postPtr++;
			}
			rowSize[i] = dataPtr - rowPtrs[i];
		}
	}

	/////////////////////////////////////////////////////////////////
	// SparseMatrix::GetRowPtr()
	//
	// Returns the pointer to a particular row in the sparse matrix.
	/////////////////////////////////////////////////////////////////

	SafeVector<PIF>::iterator GetRowPtr(int row) const {
		assert(row >= 1 && row <= seq1Length);
		return rowPtrs[row];
	}

	/////////////////////////////////////////////////////////////////
	// SparseMatrix::GetValue()
	//
	// Returns value at a particular row, column.
	/////////////////////////////////////////////////////////////////

	float GetValue(int row, int col) {
		assert(row >= 1 && row <= seq1Length);
		assert(col >= 1 && col <= seq2Length);
		for (int i = 0; i < rowSize[row]; i++) {
			if (rowPtrs[row][i].first == col)
				return rowPtrs[row][i].second;
		}
		return 0;
	}

	/////////////////////////////////////////////////////////////////
	// SparseMatrix::GetRowSize()
	//
	// Returns the number of entries in a particular row.
	/////////////////////////////////////////////////////////////////

	int GetRowSize(int row) const {
		assert(row >= 1 && row <= seq1Length);
		return rowSize[row];
	}

	/////////////////////////////////////////////////////////////////
	// SparseMatrix::GetSeq1Length()
	//
	// Returns the first dimension of the matrix.
	/////////////////////////////////////////////////////////////////

	int GetSeq1Length() const {
		return seq1Length;
	}

	/////////////////////////////////////////////////////////////////
	// SparseMatrix::GetSeq2Length()
	//
	// Returns the second dimension of the matrix.
	/////////////////////////////////////////////////////////////////

	int GetSeq2Length() const {
		return seq2Length;
	}

	/////////////////////////////////////////////////////////////////
	// SparseMatrix::GetRowPtr
	//
	// Returns the pointer to a particular row in the sparse matrix.
	/////////////////////////////////////////////////////////////////

	int GetNumCells() const {
		return data.size();
	}

	/////////////////////////////////////////////////////////////////
	// SparseMatrix::Print()
	//
	// Prints out a sparse matrix.
	/////////////////////////////////////////////////////////////////

	void Print(ostream &outfile) const {
		outfile << "Sparse Matrix:" << endl;
		for (int i = 1; i <= seq1Length; i++) {
			outfile << "  " << i << ":";
			for (int j = 0; j < rowSize[i]; j++) {
				outfile << " (" << rowPtrs[i][j].first << ","
						<< rowPtrs[i][j].second << ")";
			}
			outfile << endl;
		}
	}

	void Print() const {
		cout << "Sparse Matrix:" << endl;
		for (int i = 1; i <= seq1Length; i++) {
			cout << "  " << i << ":";
			for (int j = 0; j < rowSize[i]; j++) {
				if(rowPtrs[i][j].second > 0) cout << " (" << rowPtrs[i][j].first << ","
						<< rowPtrs[i][j].second << ")";
			}
			cout << endl;
		}
	}

	/////////////////////////////////////////////////////////////////
	// SparseMatrix::ComputeTranspose()
	//
	// Returns a new sparse matrix containing the transpose of the
	// current matrix.
	/////////////////////////////////////////////////////////////////

	SparseMatrix *ComputeTranspose() const {

		// create a new sparse matrix
		SparseMatrix *ret = new SparseMatrix();
		int numCells = data.size();

		ret->seq1Length = seq2Length;
		ret->seq2Length = seq1Length;

		// allocate memory
		ret->data.resize(numCells);
		ret->rowSize.resize(seq2Length + 1);
		ret->rowSize[0] = -1;
		ret->rowPtrs.resize(seq2Length + 1);
		ret->rowPtrs[0] = ret->data.end();

		// compute row sizes
		for (int i = 1; i <= seq2Length; i++)
			ret->rowSize[i] = 0;
		for (int i = 0; i < numCells; i++)
			ret->rowSize[data[i].first]++;

		// compute row ptrs
		for (int i = 1; i <= seq2Length; i++) {
			ret->rowPtrs[i] =
					(i == 1) ?
							ret->data.begin() :
							ret->rowPtrs[i - 1] + ret->rowSize[i - 1];
		}

		// now fill in data
		SafeVector<SafeVector<PIF>::iterator> currPtrs = ret->rowPtrs;

		for (int i = 1; i <= seq1Length; i++) {
			SafeVector<PIF>::iterator row = rowPtrs[i];
			for (int j = 0; j < rowSize[i]; j++) {
				currPtrs[row[j].first]->first = i;
				currPtrs[row[j].first]->second = row[j].second;
				currPtrs[row[j].first]++;
			}
		}

		return ret;
	}

	/////////////////////////////////////////////////////////////////
	// SparseMatrix::GetPosterior()
	//
	// Return the posterior representation of the sparse matrix.
	/////////////////////////////////////////////////////////////////

	VF *GetPosterior() const {

		// create a new posterior matrix
		VF *posteriorPtr = new VF((seq1Length + 1) * (seq2Length + 1));
		assert(posteriorPtr);
		VF &posterior = *posteriorPtr;

		// build the posterior matrix
		for (int i = 0; i < (seq1Length + 1) * (seq2Length + 1); i++)
			posterior[i] = 0;
		for (int i = 1; i <= seq1Length; i++) {
			VF::iterator postPtr = posterior.begin() + i * (seq2Length + 1);
			for (int j = 0; j < rowSize[i]; j++) {
				postPtr[rowPtrs[i][j].first] = rowPtrs[i][j].second;
			}
		}

		return posteriorPtr;
	}

};

#endif
